import numpy as np
import matplotlib.pyplot as plt
import os
import probml_utils as pml  # pip install git+https://github.com/probml/probml-utils.git
import seaborn as sns
import pandas as pd
import sklearn
from sklearn.datasets import load_iris


if __name__ == '__main__':
    sns.set(style="ticks", color_codes=True, rc={
        'figure.figsize': (4, 4),
    })

    pd.set_option('display.precision', 2)  # 2 decimal places
    pd.set_option('display.max_rows', 20)
    pd.set_option('display.max_columns', 30)
    pd.set_option('display.width', 100)  # wide windows

    iris = load_iris()

    # Extract numpy arrays
    X = iris.data
    y = iris.target

    # Convert to pandas dataframe
    df = pd.DataFrame(data=X, columns=iris.feature_names)
    df['label'] = pd.Series(iris.target_names[y], dtype='category')

    # we pick a color map to match that used by decision tree graphviz
    # cmap = ListedColormap(['#fafab0','#a0faa0', '#9898ff']) # orange, green, blue/purple
    # cmap = ListedColormap(['orange', 'green', 'purple'])
    palette = {'setosa': 'orange', 'versicolor': 'green', 'virginica': 'purple'}

    g = sns.pairplot(df, vars=df.columns[0:4], hue="label", palette=palette)
    # g = sns.pairplot(df, vars = df.columns[0:4], hue="label")
    os.environ["FIG_DIR"] = '.'
    pml.savefig('iris_scatterplot_purple.tmp.pdf')
    plt.show()

